/*
* Copyright (C) 2010-2012 Magnus Skjegstad
* 
* This program is free software; you can redistribute it and/or
* modify it under the terms of the GNU General Public License
* as published by the Free Software Foundation; either version 2
* of the License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program; if not, write to the Free Software
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
* 
* The latest version of this library can be downloaded from http://mist-pubsub.googlecode.com
* 
*/

package com.skjegstad.mist2.modules.tcp;

import com.skjegstad.simpleevents.SimpleEvent;
import com.skjegstad.simpleevents.SimpleEventListener;
import com.skjegstad.mist2.IMistNetworkModule;
import com.skjegstad.mist2.NetworkModuleHelper;
import com.skjegstad.mist2.datatypes.BloomfilterAcknowledgment;
import com.skjegstad.mist2.datatypes.ConnectionIdentifier;
import com.skjegstad.mist2.datatypes.DataFragment;
import com.skjegstad.mist2.datatypes.NodeIdentifier;
import com.skjegstad.mist2.datatypes.SimpleEntry;
import com.skjegstad.mist2.datatypes.TopicIdentifier;
import com.skjegstad.mist2.datatypes.TopicMap;
import com.skjegstad.mist2.datatypes.TopicMapEntry;
import com.skjegstad.mist2.events.NewDataElementReceivedEvent;
import com.skjegstad.mist2.events.NewDataFragmentsReceived;
import com.skjegstad.mist2.events.NewNodeDiscoveredEvent;
import com.skjegstad.mist2.exceptions.MistException;
import com.skjegstad.mist2.managers.MistEventManager;
import com.skjegstad.mist2.modules.udp.BeaconMessage;
import com.skjegstad.mist2.modules.udp.DataFragmentMessage;
import com.skjegstad.mist2.modules.udp.MessageContainer;
import java.io.EOFException;
import java.io.IOException;
import java.io.ObjectInput;
import java.io.ObjectInputStream;
import java.io.ObjectOutput;
import java.io.ObjectOutputStream;
import java.net.InetAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Random;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

/**
 *
 * @author Magnus Skjegstad
 */
public class TCPNetworkModule implements IMistNetworkModule {
    private Random r = new Random();
    private int listenPort;
    private ServerSocket serverSocket;
    private ExecutorService executor;
    private NetworkModuleHelper networkModuleHelper;
    private MistEventManager eventManager;    
    private final ConcurrentMap<Socket, Entry<ObjectInput, ObjectOutput>> connectionStreams;
    private final ConcurrentMap<Socket, NodeIdentifier> connectionNodeIdentifiers;
    private final List<Socket> connectionSockets;
    private int fragmentSize;
    private ConnectionIdentifier connectionIdentifier;
    private NodeIdentifier nodeIdentifier;
    private int hashesInLastAck = -1;
    private byte extraBits = 14;
    private int k = 1;
    private long lastBeaconSent = 0;
    private final Map<NodeIdentifier, Long> lastSentFragments = new HashMap();
    private boolean shouldSendBeacon = true;
    private int periodicBeaconTime = 60000;
    private int timerTick = 250;
    private BeaconMessage lastBeaconMessage;

    public TCPNetworkModule(int listenPort) {
        this.listenPort = listenPort;
        fragmentSize = 64000;
        this.connectionStreams = new ConcurrentHashMap();
        this.connectionNodeIdentifiers = new ConcurrentHashMap();
        this.connectionSockets = Collections.synchronizedList(new ArrayList());
    }

    private boolean addSocket(final Socket socket) {
        InetAddress src = socket.getInetAddress();        

        try {
            final ObjectOutputStream outStream = new ObjectOutputStream((socket.getOutputStream()));
            final ObjectInputStream inStream = new ObjectInputStream((socket.getInputStream()));

            connectionSockets.add(socket);
            connectionStreams.put(socket, new SimpleEntry(inStream, outStream));
            //connectionNodeIdentifiers.put(socket, null);
            
            networkModuleHelper.debug("Accepted connection from " + src);
            // send current beacon to new connection
            if (lastBeaconMessage == null)                
                sendBeacon(createBeacon(), outStream);
            else
                sendBeacon(lastBeaconMessage, outStream);

            // add reader thread for the connection
            executor.submit(new Runnable() {
                public void run() {
                    ObjectInputStream ois = inStream;
                    try {
                        networkModuleHelper.debug("In listen thread");
                        while (!Thread.interrupted() &&
                                socket.isConnected() &&
                                !socket.isInputShutdown()) {
                            try {
                                MessageContainer mc = new MessageContainer();
                                networkModuleHelper.debug("Waiting...");
                                mc = (MessageContainer)ois.readObject();
                                networkModuleHelper.debug("Got packet");
                                
                                if (mc.getMessageObject() instanceof BeaconMessage) {
                                    BeaconMessage bm = (BeaconMessage)mc.getMessageObject();
                                    readBeaconMessage(socket, bm);
                                } else
                                if (mc.getMessageObject() instanceof DataFragmentMessage) {
                                    DataFragmentMessage dfm = (DataFragmentMessage)mc.getMessageObject();
                                    networkModuleHelper.debug("Got data fragment");

                                    // skip our own fragments
                                    if (dfm.getFragment().getSourceNode().equals(networkModuleHelper.getMyNodeIdentifier())) {
                                        networkModuleHelper.debug("Got own fragment! Should not happen... Sender: " + connectionNodeIdentifiers.get(socket));
                                        continue;
                                    }

                                    networkModuleHelper.addDataFragments(Collections.singletonList(dfm.getFragment()));
                                    shouldSendBeacon = true;
                                    
                                } else
                                    networkModuleHelper.debug("Received unknown object on stream");
                                
                            } catch (EOFException ex) {
                                networkModuleHelper.debug("Remotely disconnected from " + socket.getInetAddress());
                                eventManager.throwException(new MistException("Got EOFException", ex));
                                break;
                            } catch (Exception ex) {
                                eventManager.throwException(new MistException("Got exception while reading from socket. Retrying in 1 sec.", ex));
                                try { Thread.sleep(1000); } catch (InterruptedException ex1) {}
                                continue;
                            }
                        }
                    } catch (Exception ex) {
                        eventManager.throwException(new MistException("Got exception in listen thread. Exiting...", ex));
                    } finally {
                        NodeIdentifier n = connectionNodeIdentifiers.get(socket);
                        networkModuleHelper.debug("About to exit listen thread on node " + socket.getInetAddress() +
                                ", node=" + n +
                                ", connected=" + socket.isConnected() +
                                ", inputShutdown=" + socket.isInputShutdown() +
                                ", outputShutdown=" + socket.isOutputShutdown() +
                                ", interrupted=" + Thread.interrupted());
                        removeAndCloseSocket(socket);
                    }
                }
            });

            return true;
        } catch (Exception ex) {
            connectionStreams.remove(socket);
            connectionNodeIdentifiers.remove(socket);
            connectionSockets.remove(socket);

            eventManager.throwException(new MistException("Unable to accept connection from " + src,  ex));
            networkModuleHelper.log("Refused connection from " + src);
            return false;
        }
    }

    public boolean connect(InetAddress dst, int port) {
        networkModuleHelper.debug("Connecting to " + dst + ":" + port);
        try {
            Socket s = new Socket(dst, port);
            return addSocket(s);
        } catch (Exception ex) {
            eventManager.throwException(new MistException("Unable to connect to " + dst,  ex));
            return false;
        }
    }

    private void removeAndCloseSocket(Socket socket) {
        try {
            if (socket.isConnected())
                socket.close();
        } catch (Exception ex) {
            eventManager.throwException(new MistException("Unable to terminate connection from " + socket.getInetAddress(), ex));
        } finally {
            NodeIdentifier n = connectionNodeIdentifiers.get(socket);

            connectionSockets.remove(socket);
            connectionNodeIdentifiers.remove(socket);
            connectionStreams.remove(socket);
            
            if (n != null) 
                networkModuleHelper.removeNeighborNode(n);
        }
    }

    public void start(final NetworkModuleHelper networkModuleHelper) throws MistException {
        this.networkModuleHelper = networkModuleHelper;
        this.eventManager = networkModuleHelper.getMistEventManager();
        lastSentFragments.clear();
        shouldSendBeacon = true;

        try {
            serverSocket = new ServerSocket(listenPort, 10, InetAddress.getByName("0.0.0.0"));
        } catch (IOException ex) {
            throw new MistException("Unable to open server socket on port " + listenPort, ex);
        }
        
        executor = Executors.newCachedThreadPool();

        // Add socket listening thread
        executor.submit(new Runnable() {
            public void run() {
                    while (!Thread.interrupted()) {
                        final Socket connectionSocket;
                        try {
                            connectionSocket = serverSocket.accept();
                            addSocket(connectionSocket);                          
                        } catch (Exception ex) {
                            eventManager.throwException(new MistException("Unable to accept connection on socket. Retrying in 10 sec.", ex));
                            try { Thread.sleep(10000); } catch (InterruptedException ex1) {}
                            continue;
                        }
                }
            }
        });

        // Periodic beacon thread
        executor.submit(new Runnable() {
            public void run() {
                while (!Thread.interrupted()) {
                    long timeDiff = System.currentTimeMillis() - lastBeaconSent;
                    long wait = periodicBeaconTime - timeDiff;

                    if (wait > 0)
                        try { Thread.sleep(wait + randomMs(timerTick, 0.5)); } catch (Exception ex) { continue; }
                    
                    if (timeDiff < periodicBeaconTime)
                        continue;

                    try {
                        if (lastBeaconSent <= 0)
                            networkModuleHelper.debug("--- Sending initial beacon ---");
                        else
                            networkModuleHelper.debug(("--- Sending periodic beacon after timer expired --- "));
                        sendBeaconToAll();
                        lastBeaconSent = System.currentTimeMillis(); // set this in case we didn't have any neighbors
                        shouldSendBeacon = false;
                    } catch (MistException ex) {
                        eventManager.throwException(new MistException("Unable to send periodic beacon, retrying in 10 seconds", ex));
                        try { Thread.sleep(10000); } catch (Exception ex1) {}
                        continue;
                    }
                }
            }
        });

        // Neighbor update thread
        executor.submit(new Runnable() {
            public void run() {
                while (!Thread.interrupted()) {
                    // send beacon
                    if (shouldSendBeacon) {
                        try {
                            shouldSendBeacon = false;
                            sendBeaconToAll();
                        } catch (Exception ex) {
                            eventManager.throwException(new MistException("Unable to send beacon. Retrying in 10 sec.", ex));
                            try { Thread.sleep(10000); } catch (InterruptedException ex1) {}
                            continue;
                        }
                    }

                    // send fragments
                    sendDataFragments();

                    try {
                        Thread.sleep(randomMs(timerTick, 0.5));
                    } catch (InterruptedException ex) {
                        continue;
                    }
                }
            }
        });

        // On new data fragments, send new ack
        eventManager.getDataFragmentEventPipe(connectionIdentifier).listen(new SimpleEventListener() {
            public void receiveEvent(SimpleEvent event) {
                if (event instanceof NewDataFragmentsReceived) {
                    shouldSendBeacon = true;
                }
            }
        });
        
        // On new data elements, send new ack
        eventManager.getDataManagerEventPipe().listen(new SimpleEventListener() {
            public void receiveEvent(SimpleEvent event) {
                if (event instanceof NewDataElementReceivedEvent) {
                    shouldSendBeacon = true;
                }
            }
        });

        // On new neighbors, send new ack
        eventManager.getNodeManagerEventPipe().listen(new SimpleEventListener() {
            public void receiveEvent(SimpleEvent event) {
                if (event instanceof NewNodeDiscoveredEvent) {
                    shouldSendBeacon = true;
                }
            }
        });
    }

    private void sendDataFragments() {
        List<Socket> sockets = new ArrayList();
        sockets.addAll(connectionSockets);
        synchronized (connectionStreams) {
            for (Socket s : sockets) {
                try {
                    if (!s.isConnected()) // skip disconnected sockets
                        continue;

                    NodeIdentifier node = connectionNodeIdentifiers.get(s);
                    if (node == null) // unknown node, skip
                        continue;

                    Entry<ObjectInput, ObjectOutput> e = connectionStreams.get(s);
                    ObjectOutput oos;
                    if (e != null)
                        oos = connectionStreams.get(s).getValue();
                    else
                        continue;

                    // send hashes
                    List<Long> hashes = new ArrayList(networkModuleHelper.getHashesMissingInNode(node, true, true)); // don't send expired or invalid
                    Collections.shuffle(hashes);
                    lastSentFragments.put(node, System.currentTimeMillis());
                    Collection<DataFragment> fragments = networkModuleHelper.getDataFragments(hashes, true, true);

                    for (DataFragment f : fragments) {
                        try {
                            DataFragmentMessage dfm = new DataFragmentMessage(f);
                            MessageContainer mc = new MessageContainer(dfm);
                            oos.writeObject(mc);
                        } catch (Exception ex) {
                            throw new MistException("Unable to send data fragment " + f.toString(), ex);
                        }
                    }
                    oos.flush();

                } catch (Exception ex) {
                    eventManager.throwException(new MistException("Got exception while sending data fragments to " + s.getInetAddress(), ex));
                    continue;
                }                
            }
        }
    }

    private void sendBeaconToAll() throws MistException {
        if (networkModuleHelper.getNodes().size() <= 0) { // no neighbors
            networkModuleHelper.debug("TCP: Beacon not sent. Not connected.");
            return;
        }

        BeaconMessage bm = createBeacon();
        synchronized(connectionStreams) {
            for (Entry<ObjectInput, ObjectOutput> e : connectionStreams.values()) {
                try {
                    sendBeacon(bm, e.getValue());
                } catch (MistException ex) {
                    eventManager.throwException(new MistException("Unable to send beacon on output stream/socket. Retrying later.", ex));
                    continue;
                }
            }
        }
    }

    private void sendBeacon(BeaconMessage bm, final ObjectOutput oos) throws MistException {
        try {
            lastBeaconSent = System.currentTimeMillis();
            lastBeaconMessage = bm;
            MessageContainer mc = new MessageContainer(bm);
            synchronized (connectionStreams) {
                oos.writeObject(mc);
                oos.flush();
            }
        } catch (Exception ex) {
            throw new MistException("Unable to send beacon message object on output stream", ex);
        }
    }

    private int randomMs(int ms, double perm) {
        double part = ms * perm;
        int ret = (int)Math.round(part + (r.nextDouble() * ms));
        //System.err.println(ret);
        return ret;
    }

    public void stop() throws MistException {
        executor.shutdownNow();
        try {
            executor.awaitTermination(1, TimeUnit.SECONDS);
        } catch (InterruptedException ex) {}
        
         // close all connections
        List<Socket> socketsToRemove = new ArrayList();
        synchronized(connectionSockets) {
            socketsToRemove.addAll(connectionSockets);
        }
        for (Socket s : socketsToRemove)
            removeAndCloseSocket(s);

        connectionNodeIdentifiers.clear();
        connectionSockets.clear();
        connectionStreams.clear();
        
        try {
            serverSocket.close();
        } catch (IOException ex) {
            throw new MistException("Unable to close server socket", ex);
        }
    }

    private synchronized BeaconMessage createBeacon() throws MistException {
        try {
            networkModuleHelper.debug("Creating beacon w/k=" + k);
            TopicMap topics = networkModuleHelper.getActiveTopics();
            Collection<Long> hashes = networkModuleHelper.getHashesToAck(networkModuleHelper.getActiveTopics().keySet(), false, false); // include expired and invalid in beacon until they are clean up

            if (hashes.size() == hashesInLastAck && !hashes.isEmpty())
                extraBits = (byte) ((extraBits + k) % networkModuleHelper.getMaximumExtraBits());
            else
                hashesInLastAck = hashes.size();

            if (hashes.isEmpty())
                networkModuleHelper.debug("Warning: Sending empty ack");

            BeaconMessage bm;
            bm = new BeaconMessage(networkModuleHelper.getMyNodeIdentifier(),
                topics,
                listenPort,
                networkModuleHelper.getMaximumAge(),
                new BloomfilterAcknowledgment(hashes, extraBits, k));

            return bm;
        } catch (Exception ex) {
            throw new MistException("Unable to create beacon", ex);
        }
    }

    public synchronized void readBeaconMessage(Socket srcSocket, BeaconMessage m) {
        if (!m.getNodeIdentifier().equals(networkModuleHelper.getMyNodeIdentifier())) {
            String s = "";
            if (m.getAcknowledgements() != null)
                for (int i = 0; i < m.getAcknowledgements().getData().length; i++) {
                    if (!s.isEmpty())
                        s = s + ".";
                    s = s + m.getAcknowledgements().getData()[i];
                }

            // store node identifier
            connectionNodeIdentifiers.put(srcSocket, m.getNodeIdentifier());

            networkModuleHelper.debug("Got beacon from " + m.getNodeIdentifier() + "@" + srcSocket.getInetAddress() + ":" + m.getContactPort() +
                    ", len: " + m.getAcknowledgements().getLengthInBits() +
                    " bytes: " + m.getAcknowledgements().getData().length  +
                    ", ackdata: \"" + s + "\"");

            for (Entry<TopicIdentifier, TopicMapEntry> et : m.getSubscribeToTopics().entrySet())
                    networkModuleHelper.debug("Got topic in beacon " + et.getKey().getHashes() + ", ttl:" + et.getValue().getTtl() + ", checksum: " + et.getValue().getChecksum());
            
            lastSentFragments.remove(m.getNodeIdentifier());

            networkModuleHelper.registerNeighborNode(m.getNodeIdentifier(),
                    m.getSubscribeToTopics(),
                    srcSocket.getInetAddress(),
                    m.getAcknowledgements());
        }
    }

    public int getFragmentSize() {
        return fragmentSize;
    }

    public void setConnectionIdentifier(ConnectionIdentifier connectionIdentifier) {
        this.connectionIdentifier = connectionIdentifier;
    }

    public ConnectionIdentifier getConnectionIdentifier() {
        return this.connectionIdentifier;
    }

    public void setNodeIdentifier(NodeIdentifier nodeIdentifier) {
        this.nodeIdentifier = nodeIdentifier;
    }

    public NodeIdentifier getNodeIdentifier() {
        return this.nodeIdentifier;
    }

}
